{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 总结：Quantum Variational Autoencoder（QVAE）原理与代码实现\n",
    "\n",
    "---\n",
    "\n",
    "## 一、QVAE 原理概括\n",
    "\n",
    "QVAE（Quantum Variational Autoencoder）是一种将 **量子生成模型** 引入 **变分自编码器(VAE)** 潜空间的生成模型。其核心思想是：\n",
    "\n",
    "> **用量子玻尔兹曼机（QBM）替代传统VAE中的先验分布，从而构建一个具有量子生成能力的潜变量模型。**\n",
    "\n",
    "### 模型结构\n",
    "\n",
    "QVAE 包括以下关键组件：\n",
    "\n",
    "1. **编码器（Encoder）**：  \n",
    "   将输入数据 x 映射为潜变量的近似后验分布  \n",
    "   $q_\\phi(z|x)$，通常由神经网络参数化。\n",
    "\n",
    "2. **先验分布（Prior）**：  \n",
    "   使用 **量子玻尔兹曼机(QBM)** 建模潜变量 z 的先验分布。哈密顿量为：  \n",
    "   $$\n",
    "   \\mathcal{H}_\\theta = \\sum_l \\Gamma_l \\sigma_l^x + \\sum_l h_l \\sigma_l^z + \\sum_{l<m} W_{lm} \\sigma_l^z \\sigma_m^z\n",
    "   $$\n",
    "\n",
    "3. **解码器（Decoder）**：  \n",
    "   将潜变量 z（或其连续松弛变量 $\\zeta$）映射回数据空间，并使用解码器重建原始数据：  \n",
    "   $$\n",
    "   p_\\theta(\\mathbf{x} | \\boldsymbol{\\zeta}) \\approx \\text{Bernoulli}(f_\\theta(\\boldsymbol{\\zeta}))\n",
    "   $$\n",
    "\n",
    "### 训练目标：Q-ELBO\n",
    "\n",
    "QVAE 使用一个 **量子下界(Q-ELBO)** 来近似最大化对数似然：\n",
    "\n",
    "$$\n",
    "\\mathcal{L}_{\\text{Q-ELBO}} = \\mathbb{E}_{q_\\phi(\\mathbf{z}|\\mathbf{x})} [\\log p_\\theta(\\mathbf{x} | \\boldsymbol{\\zeta})] - \\tilde{H}(q_\\phi(\\mathbf{z}|\\mathbf{x}) \\| p_\\theta(\\mathbf{z}))\n",
    "$$\n",
    "\n",
    "### QBM采样与训练\n",
    "\n",
    "- **正相（positive phase）**：从编码器采样  $z \\sim q_\\phi(z|x)$\n",
    "- **负相（negative phase）**：从 QBM 中采样 $z \\sim p_\\theta(z)$，使用 **蒙特卡洛方法** 或 **量子退火器**\n",
    "把能量作为目标函数，objective的梯度即为基于正相和负相采样计算的梯度\n",
    "\n",
    "\n",
    "## 二、代码内容概括\n",
    "\n",
    "该例子实现了一个 **基于 MNIST 的 QVAE 训练流程**，包括以下模块：\n",
    "\n",
    "\n",
    "### 模型构建\n",
    "\n",
    "- 使用自定义 `QVAE` 类（未在代码中展开，假设已封装）：\n",
    "  - 编码器/解码器为全连接网络；\n",
    "  - 使用 `mean_x` 作为模型偏移量；\n",
    "  - 包含玻尔兹曼机作为先验分布\n",
    "\n",
    "### 训练过程\n",
    "\n",
    "- 优化器：Adam，学习率 1e-3；\n",
    "- 损失函数：`loss = neg_elbo + wd_loss`，即包含ELBO下界和weight decay\n",
    "- 每个 epoch 保存模型权重；\n",
    "- 记录并保存以下历史数据：\n",
    "  - `loss_history`\n",
    "  - `elbo_history`\n",
    "  - `kl_history`\n",
    "  - `cost_history`\n",
    "\n",
    "\n",
    "### 图像生成流程\n",
    "\n",
    "1. 从 RBM 中采样潜变量 `z`；\n",
    "2. 引入平滑分布 `Exponential(beta)` 采样 `zeta`；\n",
    "3. 使用解码器生成图像：\n",
    "   $$\n",
    "   \\text{generated\\_x} = \\sigma(\\text{decoder}(zeta) + \\text{train\\_bias})\n",
    "   $$\n",
    "\n",
    "### FID 评估\n",
    "\n",
    "- 使用 `torchmetrics.image.fid.FrechetInceptionDistance`；\n",
    "- 提供函数：\n",
    "  - `get_real_images()`：从验证集获取真实图像；\n",
    "  - `generate_images_qvae()`：使用 QVAE 生成图像；\n",
    "  - `compute_fid_in_batches()`：计算 FID 分数（适配 MNIST）；\n",
    "- 最终输出 QVAE 在 MNIST 上的 FID 分数。\n",
    "\n",
    "\n",
    "\n",
    "## 三、总结\n",
    "\n",
    "> QVAE = VAE 的编码器/解码器 + QBM 的量子先验，通过最大化量子下界（Q-ELBO）实现端到端训练；  \n",
    "> 代码包含MNIST 数据集，QVAE 模型，训练循环，图像可视化和FID 评估，完整实现了量子自编码器模型的实验流程。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "import matplotlib.pyplot as plt\n",
    "from torch.utils.data import DataLoader, Dataset\n",
    "from torchvision import datasets, transforms\n",
    "from tqdm import tqdm\n",
    "\n",
    "# 添加kaiwu的license信息\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MNIST数据集加载"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MNISTWithBatch(Dataset):\n",
    "    def __init__(self, root, train=True, transform=None, download=True, num_batches=6):\n",
    "        self.mnist = datasets.MNIST(root=root, train=train, transform=transform, download=download)\n",
    "        self.num_batches = num_batches\n",
    "        self.batch_indices = self._create_batch_indices()\n",
    "\n",
    "    def _create_batch_indices(self):\n",
    "        num_samples = len(self.mnist)\n",
    "        return torch.arange(num_samples) // (num_samples // self.num_batches)\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.mnist)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        image, _ = self.mnist[idx]  # 忽略原始 label\n",
    "        batch_idx = self.batch_indices[idx]\n",
    "        return image, batch_idx"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Binarized MNIST数据集加载"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 定义数据转换操作\n",
    "def flatten_tensor(x):\n",
    "    return x.view(-1)\n",
    "transform = transforms.Compose([\n",
    "    transforms.ToTensor(),  # 将图像转换为Tensor\n",
    "    transforms.Lambda(flatten_tensor)\n",
    "])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size=256\n",
    "epochs=20\n",
    "lr=1e-3\n",
    "kl_beta = 0.000001\n",
    "\n",
    "model_name = \"QVAE_annealing_tanh\"\n",
    "save_path = f\"./models/{model_name}\"\n",
    "os.makedirs(save_path, exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "val_dataset = MNISTWithBatch(root='../../data', train=False, download=False, transform=transform)\n",
    "val_dataloader = DataLoader(val_dataset, batch_size=batch_size, shuffle=True)\n",
    "\n",
    "train_dataset = MNISTWithBatch(root='../../data', train=True, download=False, transform=transform)\n",
    "train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)\n",
    "\n",
    "mean_x = 0\n",
    "for x, _ in train_dataloader:\n",
    "    mean_x += x.mean(dim=0)\n",
    "mean_x = mean_x / len(train_dataloader)\n",
    "mean_x = mean_x.cpu().numpy()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from kaiwu.classical import SimulatedAnnealingOptimizer\n",
    "from kaiwu.torch_plugin import RestrictedBoltzmannMachine, QVAE\n",
    "from qvae_enc_dec import Encoder, Decoder\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "# 图片拉伸后的维度\n",
    "input_dim = 784\n",
    "# fc1压缩后的维度\n",
    "hidden_dim = 512\n",
    "# 隐变量维度\n",
    "latent_dim = 256\n",
    "\n",
    "# RBM可见层和隐藏层维度\n",
    "num_var1 = 128\n",
    "num_var2 = 128\n",
    "# 重叠分布的beta\n",
    "dist_beta = 10\n",
    "\n",
    "encoder = Encoder(input_dim, hidden_dim, latent_dim, weight_decay=0.01)\n",
    "rbm = RestrictedBoltzmannMachine(\n",
    "    num_visible=num_var1, num_hidden=num_var2, h_range=[-1, 1], j_range=[-1, 1]\n",
    ")\n",
    "decoder = Decoder(latent_dim, hidden_dim, input_dim, weight_decay=0.01)\n",
    "sampler = SimulatedAnnealingOptimizer(alpha=0.95)\n",
    "\n",
    "model = QVAE(encoder,decoder, rbm, sampler, dist_beta, mean_x, num_var1)\n",
    "model.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = torch.optim.Adam(model.parameters(), lr=lr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "patience_counter = 0\n",
    "best_state_dict = None\n",
    "loss_history = []\n",
    "elbo_history = []\n",
    "kl_history = []\n",
    "cost_history = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for epoch in range(1, epochs):  # 遍历每个训练轮次\n",
    "    model.train()  # 设置模型为训练模式\n",
    "    total_loss, total_elbo, total_kl, total_cost = 0, 0, 0, 0  # 初始化本轮累计损失等指标\n",
    "    for x, _ in train_dataloader:  # 遍历训练集的每个批次\n",
    "        x = x.to(device)  # 将输入数据移动到指定设备（CPU/GPU）\n",
    "\n",
    "        optimizer.zero_grad()  # 梯度清零\n",
    "\n",
    "        # 前向传播，计算负ELBO、权重衰减、KL散度等\n",
    "        output, recon_x, neg_elbo, wd_loss, kl, cost, _, _ = model.neg_elbo(x, kl_beta)\n",
    "        loss = neg_elbo + wd_loss  # 总损失 = 负ELBO + 权重衰减\n",
    "        loss.backward()  # 反向传播，计算梯度\n",
    "\n",
    "        optimizer.step()  # 更新模型参数\n",
    "\n",
    "        # 累加本批次的各项指标\n",
    "        total_loss += loss.item()\n",
    "        total_elbo += neg_elbo.item()\n",
    "        total_kl += kl.item()\n",
    "        total_cost += cost.item()\n",
    "\n",
    "    # 计算本轮各项指标的平均值\n",
    "    avg_loss = total_loss / len(train_dataloader)\n",
    "    avg_elbo = total_elbo / len(train_dataloader)\n",
    "    avg_kl = total_kl / len(train_dataloader)\n",
    "    avg_cost = total_cost / len(train_dataloader)\n",
    "\n",
    "    # 记录历史指标\n",
    "    loss_history.append(avg_loss)\n",
    "    elbo_history.append(avg_elbo)\n",
    "    kl_history.append(avg_kl)\n",
    "    cost_history.append(avg_cost)\n",
    "\n",
    "    # 保存当前轮次的模型参数\n",
    "    model_save_path = os.path.join(save_path, f'davepp_epoch{epoch}.pth')\n",
    "    torch.save(model.state_dict(), model_save_path)\n",
    "\n",
    "    # 打印本轮训练结果\n",
    "    print(f\"Epoch {epoch}/{epochs}: Loss: {avg_loss:.4f}, elbo: {avg_elbo:.4f}, KL: {avg_kl:.4f}, Cost: {avg_cost:.4f}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def save_list_to_txt(filename, data):\n",
    "    with open(filename, \"w\") as f:\n",
    "        for value in data:\n",
    "            f.write(f\"{value:.6f}\\n\")\n",
    "\n",
    "\n",
    "save_list_to_txt(os.path.join(save_path, \"loss_history.txt\"), loss_history)\n",
    "save_list_to_txt(os.path.join(save_path, \"elbo_history.txt\"), elbo_history)\n",
    "save_list_to_txt(os.path.join(save_path, \"cost_history.txt\"), cost_history)\n",
    "save_list_to_txt(os.path.join(save_path, \"kl_history.txt\"), kl_history)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_flattened_images_grid(features: torch.Tensor, grid_size: int = 8, save_path: str = None):\n",
    "    \"\"\"\n",
    "    显示并可选保存前 grid_size * grid_size 个 28x28 灰度图像。\n",
    "\n",
    "    Args:\n",
    "        features (torch.Tensor): 形状为 [N, 784] 的张量，每行为一个扁平的 28x28 图像。\n",
    "        grid_size (int): 图像网格边长（默认 8，即显示前 64 张图像）。\n",
    "        save_path (str): 如果提供，将保存图像到该路径。\n",
    "    \"\"\"\n",
    "    assert features.dim() == 2 and features.size(1) == 784, \"features 应为 [N, 784] 的张量\"\n",
    "    num_images = grid_size * grid_size\n",
    "    assert features.size(0) >= num_images, f\"features 中至少应包含 {num_images} 张图像\"\n",
    "\n",
    "    features_numpy = features[:num_images].detach().cpu().numpy()\n",
    "\n",
    "    fig, axes = plt.subplots(grid_size, grid_size, figsize=(5, 5))\n",
    "    for i in range(grid_size):\n",
    "        for j in range(grid_size):\n",
    "            idx = i * grid_size + j\n",
    "            img = features_numpy[idx].reshape(28, 28)\n",
    "            axes[i, j].imshow(img, cmap='gray')\n",
    "            axes[i, j].axis('off')\n",
    "\n",
    "    plt.tight_layout()\n",
    "    \n",
    "    if save_path:\n",
    "        plt.show()\n",
    "        plt.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 从训练数据加载器中取出一个批次的特征（features），忽略标签\n",
    "features, _ = next(iter(train_dataloader))\n",
    "\n",
    "# 将原始输入特征（展平的图像）以8x8网格形式可视化并保存为original.png\n",
    "plot_flattened_images_grid(features, grid_size=8, save_path = save_path + f'/original.png')\n",
    "\n",
    "# 将特征数据移动到指定设备（CPU或GPU）\n",
    "features = features.to(device)\n",
    "\n",
    "# 设置模型为评估模式（关闭dropout、BN等训练特性）\n",
    "model.eval()\n",
    "\n",
    "# 前向传播：用QVAE模型计算重构结果和各项损失\n",
    "output, recon_x, neg_elbo, wd_loss, kl, cost, q, zeta = model.neg_elbo(features, kl_beta)\n",
    "\n",
    "# 将重构后的图像以8x8网格形式可视化并保存为recon_x.png\n",
    "plot_flattened_images_grid(output, grid_size=8, save_path = save_path + f'/recon_x.png')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sampler = SimulatedAnnealingOptimizer(size_limit=100, alpha=0.99)\n",
    "z = model.bm.sample(sampler)\n",
    "shape = z.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.distributions import Exponential\n",
    "\n",
    "smoothing_dist = Exponential(dist_beta)\n",
    "# 从平滑分布采样\n",
    "zeta = smoothing_dist.sample(shape)\n",
    "zeta = zeta.to(z.device)\n",
    "zeta = torch.where(z == 0., zeta, 1-zeta) # 引入z\n",
    "# zeta = torch.randn(256, 256).to(device)\n",
    "generated_x = model.decoder(zeta)\n",
    "generated_x = generated_x + model.train_bias\n",
    "\n",
    "generated_x = torch.sigmoid(generated_x)\n",
    "\n",
    "plot_flattened_images_grid(generated_x, grid_size=8, save_path = save_path + f'/generated_x.png')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchmetrics.image.fid import FrechetInceptionDistance\n",
    "\n",
    "def get_real_images(dataloader, n_images=10000):\n",
    "    images = []\n",
    "    for batch_imgs, _ in dataloader:\n",
    "        images.append(batch_imgs)\n",
    "        if sum(img.shape[0] for img in images) >= n_images:\n",
    "            break\n",
    "    return torch.cat(images, dim=0)[:n_images]\n",
    "\n",
    "def generate_images_original_vae(model, latent_dim, n_images=10000, batch_size=64):\n",
    "    model.eval()\n",
    "    imgs = []\n",
    "    with torch.no_grad():\n",
    "        for _ in tqdm(range(n_images // batch_size)):\n",
    "            z = torch.randn(batch_size, latent_dim).to(device)\n",
    "            img = model.decoder(z).cpu()\n",
    "            imgs.append(img)\n",
    "    return torch.cat(imgs, dim=0)[:n_images]\n",
    "\n",
    "def generate_images_qvae(model, latent_dim, n_images=10000, batch_size=64):\n",
    "    model.eval()\n",
    "    imgs = []\n",
    "    sampler = SimulatedAnnealingOptimizer(alpha=0.95)\n",
    "    with torch.no_grad():\n",
    "        for _ in tqdm(range(n_images // batch_size)):\n",
    "            z = model.bm.sample(sampler)\n",
    "            shape = z.shape\n",
    "            smoothing_dist = Exponential(dist_beta)\n",
    "            # 从平滑分布采样\n",
    "            zeta = smoothing_dist.sample(shape)\n",
    "            zeta = zeta.to(z.device)\n",
    "            zeta = torch.where(z == 0., zeta, 0)\n",
    "            # zeta = torch.randn(256, 256).to(device)\n",
    "            generated_x = model.decoder(zeta)\n",
    "            \n",
    "            generated_x = generated_x + model.train_bias\n",
    "\n",
    "            generated_x = torch.sigmoid(generated_x)\n",
    "            \n",
    "            imgs.append(generated_x)\n",
    "    return torch.cat(imgs, dim=0)[:n_images]\n",
    "\n",
    "resize = transforms.Resize((299, 299))\n",
    "\n",
    "def preprocess(images):\n",
    "    if images.max() > 1.0:\n",
    "        images = images / 255.0\n",
    "    return resize(images)\n",
    "\n",
    "\n",
    "def compute_fid_in_batches(fake_imgs, real_imgs, batch_size=64):\n",
    "    \"\"\"\n",
    "    计算 FID 分数，适用于输入为 (N, 784) 的展平图像（如 MNIST）\n",
    "    \n",
    "    参数:\n",
    "        fake_imgs: 生成图像，shape = (N, 784)\n",
    "        real_imgs: 真实图像，shape = (M, 784)\n",
    "        batch_size: 每个批次处理多少图像\n",
    "        device: 使用 'cuda' 或 'cpu'\n",
    "    返回:\n",
    "        FID 分数\n",
    "    \"\"\"\n",
    "    fid = FrechetInceptionDistance(feature=64).to(device)\n",
    "\n",
    "    def preprocess(images):\n",
    "        # 转换为图像格式 (B, 1, 28, 28)\n",
    "        images = images.view(-1, 1, 28, 28)\n",
    "        # 扩展为三通道\n",
    "        images = images.repeat(1, 3, 1, 1)\n",
    "        # 调整大小到 299x299\n",
    "        resize = transforms.Resize((299, 299), antialias=True)\n",
    "        return resize(images)\n",
    "\n",
    "    # 如果不是 tensor，先转成 tensor\n",
    "    if not isinstance(fake_imgs, torch.Tensor):\n",
    "        fake_imgs = torch.tensor(fake_imgs, dtype=torch.uint8)\n",
    "    if not isinstance(real_imgs, torch.Tensor):\n",
    "        real_imgs = torch.tensor(real_imgs, dtype=torch.uint8)\n",
    "\n",
    "    # 归一化到 [0, 255] 并转为 uint8（假定输入是 float 在 [0,1] 范围）\n",
    "    fake_imgs = (fake_imgs * 255).clamp(0, 255).to(torch.uint8)\n",
    "    real_imgs = (real_imgs * 255).clamp(0, 255).to(torch.uint8)\n",
    "\n",
    "    # 转换为图像并更新 FID\n",
    "    for i in range(0, len(real_imgs), batch_size):\n",
    "        batch = real_imgs[i:i+batch_size]\n",
    "        batch = preprocess(batch)\n",
    "        fid.update(batch.to(device), real=True)\n",
    "\n",
    "    for i in range(0, len(fake_imgs), batch_size):\n",
    "        batch = fake_imgs[i:i+batch_size]\n",
    "        batch = preprocess(batch)\n",
    "        fid.update(batch.to(device), real=False)\n",
    "\n",
    "    return fid.compute().item()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 获取真实图像和生成图像\n",
    "real_imgs = get_real_images(val_dataloader, n_images=10000)\n",
    "print(f\"Real images shape: {real_imgs.shape}\")\n",
    "\n",
    "fake_imgs_original_vae = generate_images_qvae(model, latent_dim=latent_dim)\n",
    "print(f\"Generated images shape: {fake_imgs_original_vae.shape}\")\n",
    "\n",
    "# 计算 FID（更节省内存）\n",
    "fid_original = compute_fid_in_batches(fake_imgs_original_vae, real_imgs)\n",
    "print(f\"Original VAE FID: {fid_original:.2f}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py310",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
